﻿using System;
using System.Collections.Generic;
using System.Linq;

namespace LinqExtensions
{
    /// <summary>
    /// Variant of implementation of Knuth–Morris–Pratt string search algorithm (KMP algorithm).
    /// see http://en.wikipedia.org/wiki/Knuth%E2%80%93Morris%E2%80%93Pratt_algorithm
    /// </summary>
    public static class KnuthMorrisPratt
    {
        /// <summary>
        /// Reports the zero-based index of the first occurrence of the specified subsequence in given sequence.
        /// </summary>
        /// <remarks>
        /// Function implements variant of Knuth–Morris–Pratt string search algorithm (KMP algorithm).
        /// see http://en.wikipedia.org/wiki/Knuth%E2%80%93Morris%E2%80%93Pratt_algorithm
        /// </remarks>
        /// <seealso cref="KnuthMorrisPratt.IndexOfSequenceKMP{T}(IEnumerable{T}, IEnumerable{T})"/>
        /// <seealso cref="KnuthMorrisPratt.IndexOfSequenceKMP{T}(IEnumerable{T}, FailureTable{T})"/>
        /// <typeparam name="T">Element type</typeparam>
        /// <param name="source">Sequence that is searched for the specified subsequence</param>
        /// <param name="sequence">Sequence to be found</param>
        /// <param name="equalityComparer">Equality comparer</param>
        /// <returns>-1 if sequence is not found, or 0-based index of the first occurence</returns>
        public static int IndexOfSequenceKMP<T>(this IEnumerable<T> source, IEnumerable<T> sequence, IEqualityComparer<T> equalityComparer)
        {
            if (source == null) throw new ArgumentNullException("source");
            if (sequence == null) throw new ArgumentNullException("sequence");
            if (equalityComparer == null) throw new ArgumentNullException("equalityComparer");

            var indexedSequence = sequence as IList<T> ?? sequence.ToArray();
            var count = indexedSequence.Count;
            if (count == 0)
                return 0;
            if (count == 1)
                return Search.InternalIndexOf(source, indexedSequence[0], equalityComparer);

            var table = new FailureTable<T>(indexedSequence, equalityComparer);
            return InternalIndexOfSequence(source, table);
        }

        /// <summary>
        /// Reports the zero-based index of the first occurrence of the specified subsequence in given sequence.
        /// </summary>
        /// <remarks>
        /// Function implements variant of Knuth–Morris–Pratt string search algorithm (KMP algorithm).
        /// see http://en.wikipedia.org/wiki/Knuth%E2%80%93Morris%E2%80%93Pratt_algorithm
        /// </remarks>
        /// <seealso cref="KnuthMorrisPratt.IndexOfSequenceKMP{T}(IEnumerable{T}, IEnumerable{T}, IEqualityComparer{T})"/>
        /// <seealso cref="KnuthMorrisPratt.IndexOfSequenceKMP{T}(IEnumerable{T}, FailureTable{T})"/>
        /// <typeparam name="T">Element type</typeparam>
        /// <param name="source">Sequence that is searched for the specified subsequence</param>
        /// <param name="sequence">Sequence to be found</param>
        /// <returns>-1 if sequence is not found, or 0-based index of the first occurence</returns>
        public static int IndexOfSequenceKMP<T>(this IEnumerable<T> source, IEnumerable<T> sequence)
        {
            return IndexOfSequenceKMP(source, sequence, EqualityComparer<T>.Default);
        }

        /// <summary>
        /// Reports the zero-based index of the first occurrence of the specified subsequence in given sequence.
        /// </summary>
        /// <remarks>
        /// Function implements variant of Knuth–Morris–Pratt string search algorithm (KMP algorithm).
        /// see http://en.wikipedia.org/wiki/Knuth%E2%80%93Morris%E2%80%93Pratt_algorithm
        /// Subsequence to be found is in a preprocessed form, thus eliminating table building step of
        /// Knuth–Morris–Pratt algorithm on each subsequent run.
        /// </remarks>
        /// <seealso cref="KnuthMorrisPratt.IndexOfSequenceKMP{T}(IEnumerable{T}, IEnumerable{T}, IEqualityComparer{T})"/>
        /// <seealso cref="KnuthMorrisPratt.PrecomputeTable{T}(IEnumerable{T}, IEqualityComparer{T})"/>
        /// <typeparam name="T">Element type</typeparam>
        /// <param name="source">Sequence that is searched for the specified subsequence</param>
        /// <param name="failureTable">Sequence to be found, in preprocessed form. <see cref="KnuthMorrisPratt.PrecomputeTable{T}(IEnumerable{T}, IEqualityComparer{T})"/></param>
        /// <returns>-1 if sequence is not found, or 0-based index of the first occurence</returns>
        public static int IndexOfSequenceKMP<T>(this IEnumerable<T> source, FailureTable<T> failureTable)
        {
            if (source == null) throw new ArgumentNullException("source");
            if (failureTable == null) throw new ArgumentNullException("failureTable");

            if (failureTable.Sequence.Count == 0)
                return 0;
            if (failureTable.Sequence.Count == 1)
                return Search.InternalIndexOf(source, failureTable.Sequence[0], failureTable.EqualityComparer);

            return InternalIndexOfSequence(source, failureTable);
        }

        /// <summary>
        /// Precomputes 'failure table' for the given search sequence, so it can be used in KMP search
        /// </summary>
        /// <seealso cref="KnuthMorrisPratt.IndexOfSequenceKMP{T}(IEnumerable{T}, FailureTable{T})"/>
        /// <typeparam name="T">Element type</typeparam>
        /// <param name="sequence">Sequence to be found</param>
        /// <returns>Precomputed failure table</returns>
        public static FailureTable<T> PrecomputeTable<T>(IEnumerable<T> sequence)
        {
            return new FailureTable<T>(sequence, EqualityComparer<T>.Default);
        }

        /// <summary>
        /// Precomputes 'failure table' for the given search sequence, so it can be used in KMP search
        /// </summary>
        /// <seealso cref="KnuthMorrisPratt.IndexOfSequenceKMP{T}(IEnumerable{T}, FailureTable{T})"/>
        /// <typeparam name="T">Element type</typeparam>
        /// <param name="sequence">Sequence to be found</param>
        /// <param name="equalityComparer">Equality comparer</param>
        /// <returns>Precomputed failure table</returns>
        public static FailureTable<T> PrecomputeTable<T>(IEnumerable<T> sequence, IEqualityComparer<T> equalityComparer)
        {
            if (sequence == null) throw new ArgumentNullException("sequence");
            if (equalityComparer == null) throw new ArgumentNullException("equalityComparer");
            return new FailureTable<T>(sequence, equalityComparer);
        }


        private static int InternalIndexOfSequence<T>(IEnumerable<T> source, FailureTable<T> failureTable)
        {
            var equalityComparer = failureTable.EqualityComparer;
            var shifts = failureTable.Shifts;
            var sequenceArr = failureTable.Sequence;
            var sequenceLen = sequenceArr.Count;

            int sourcePosition = -1, sequencePosition = 0;
            var readNext = true;
            using (var enumerator = source.GetEnumerator())
                while (true)
                {
                    if (readNext)
                    {
                        if (!enumerator.MoveNext())
                            return -1;
                        sourcePosition++;
                    }

                    readNext = true;
                    if (equalityComparer.Equals(enumerator.Current, sequenceArr[sequencePosition]))
                    {
                        sequencePosition++;
                        if (sequencePosition >= sequenceLen)
                            return sourcePosition - sequencePosition + 1;
                    }
                    else if (sequencePosition > 0)
                    {
                        sequencePosition = shifts[sequencePosition];
                        readNext = false;
                    }
                }
        }

        /// <summary>
        /// Precomputed failure table for KMP algorithm
        /// </summary>
        /// <remarks>
        /// <see cref="KnuthMorrisPratt.PrecomputeTable{T}(IEnumerable{T}, IEqualityComparer{T})"/>
        /// </remarks>
        /// <typeparam name="T">Element type</typeparam>
        public sealed class FailureTable<T>
        {
            private readonly IEqualityComparer<T> _equalityComparer;
            private readonly IList<T> _sequence;
            private readonly int[] _shifts;

            internal FailureTable(IEnumerable<T> sequence, IEqualityComparer<T> equalityComparer)
            {
                _equalityComparer = equalityComparer;
                _sequence = sequence as IList<T> ?? sequence.ToArray();

                _shifts = new int[_sequence.Count()];
                if (_shifts.Length == 0)
                    return;

                int pos = 2, cnd = 0;
                _shifts[0] = -1;
                while (pos < _shifts.Length)
                {
                    if (equalityComparer.Equals(_sequence[pos - 1], _sequence[cnd]))
                    {
                        cnd++;
                        _shifts[pos] = cnd;
                        pos++;
                    }
                    else if (cnd > 0)
                        cnd = _shifts[cnd];
                    else pos++;
                }
            }

            internal IList<T> Sequence { get { return _sequence; } }
            internal int[] Shifts { get { return _shifts; } }
            internal IEqualityComparer<T> EqualityComparer { get { return _equalityComparer; } }
        }
    }
}
